import java.io.BufferedInputStream;
import java.io.BufferedOutputStream;
import java.io.IOException;
import java.net.ServerSocket;
import java.net.Socket;
import java.util.concurrent.CountDownLatch;

public class TCPProxy implements Runnable
{
    private String host;
    private Socket clientSoc;
    private ComputeNode cn;

    private static final int BUFFER_SIZE = 1024;

    public TCPProxy(ComputeNode cn, Socket clientSoc)
    {
        this.host = cn.hostname;
        this.cn = cn;
        this.clientSoc = clientSoc;
    }

    public void run()
    {
        try
        {
            Socket cnSoc = new Socket(host, clientSoc.getLocalPort());
            Logger.log("Connection established to " + host);
            cn.setSessions(cn.getSessions() + 1);
            CountDownLatch latch = new CountDownLatch(2);
            new Thread(new Forwarder(cnSoc, clientSoc, latch)).start();
            new Thread(new Forwarder(clientSoc, cnSoc, latch)).start();
            latch.await();
            Logger.log("Connection closed to " + host);
            cn.setSessions(cn.getSessions() - 1);
        }
        catch (IOException e)
        {
            e.printStackTrace();
        }
        catch (InterruptedException e)
        {
            // TODO Auto-generated catch block
            e.printStackTrace();
        }
    }
    
//    public static void main(String[] args) throws IOException
//    {
//        ServerSocket soc = new ServerSocket(21000);
//        while (true)
//        {
//            Socket csoc = soc.accept();
//            new TCPProxy("ec2-23-22-65-132.compute-1.amazonaws.com", csoc).run();
//        }
//    }

    private class Forwarder implements Runnable
    {
        private Socket src;
        private Socket dst;
        private CountDownLatch latch;

        public Forwarder(Socket src, Socket dst, CountDownLatch latch)
        {
            this.src = src;
            this.dst = dst;
            this.latch = latch;
        }

        @Override
        public void run()
        {
            BufferedInputStream in = null;
            BufferedOutputStream out = null;
            try
            {
                in = new BufferedInputStream(src.getInputStream());
                out = new BufferedOutputStream(dst.getOutputStream());

                byte[] buffer = new byte[BUFFER_SIZE];

                int n;
                while ((n = in.read(buffer)) != -1)
                {
                    out.write(buffer, 0, n);
                    out.flush();
                }
            }
            catch (IOException e)
            {
            }
            finally
            {
                try
                {
                    latch.countDown();
                    in.close();
                    out.close();
                    src.close();
                    dst.close();
                }
                catch (IOException e)
                {
                    e.printStackTrace();
                }
            }
        }
    }
}
